# 导入相关需要的包
import torch
import torch.nn as nn
import math

def attention(q,k,v,dropout=None):
    #将k矩阵的最后一个维度值作为d_k
    d_k = k.size(-1)
    #将k矩阵的最后两个维度互换(转置),与q矩阵相乘,除以d_k开根号
    scores = torch.matmul(q,k.transpose(-2,-1)) / math.sqrt(d_k)
    p_attn = torch.softmax(scores,dim=-1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn,v)

class SelfAtt(nn.Module):
    def __init__(self, hidden_dim):
        super(SelfAtt, self).__init__()
        self.hidden_dim = hidden_dim
        # 一般 Linear 都是默认有 bias
        # 一般来说， input dim 的 hidden dim
        self.query_proj = nn.Linear(hidden_dim, hidden_dim)
        self.key_proj = nn.Linear(hidden_dim, hidden_dim)
        self.value_proj = nn.Linear(hidden_dim, hidden_dim)

    def forward(self, x):
        # X shape is: (batch, seq_len, hidden_dim)， 一般是和 hidden_dim 相同
        # 但是 X 的 final dim 可以和 hidden_dim 不同
        q = self.query_proj(x)
        k = self.key_proj(x)
        v = self.value_proj(x)

        output = attention(q,k,v)
        return output


X = torch.rand(3, 2, 4)
net = SelfAtt(4)
y = net(X)
print(y)
print(y.shape)